// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>
#include <fstream>

#include <pollux/testing/expression/expression_verifier.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/functions/registerer.h>
#include <pollux/parse/expressions.h>
#include <pollux/parse/expressions_parser.h>
#include <pollux/parse/type_resolver.h>
#include <pollux/type/type.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::test {

namespace {

template <typename T>
struct AlwaysThrowsUserErrorFunction {
  template <typename TResult, typename TInput>
  MELON_ALWAYS_INLINE void call(TResult&, const TInput&) {
    POLLUX_USER_FAIL("expected");
  }
};

template <typename T>
struct AlwaysThrowsRuntimeErrorFunction {
  template <typename TResult, typename TInput>
  MELON_ALWAYS_INLINE void call(TResult&, const TInput&) {
    POLLUX_FAIL("expected");
  }
};

template <typename T>
struct ThrowsRuntimeErrorOnEvenInputFunction {
  template <typename TResult, typename TInput>
  MELON_ALWAYS_INLINE void call(TResult&, const TInput& in) {
    if (in % 2 == 0) {
      POLLUX_FAIL("expected");
    }
  }
};

void removeDirecrtoryIfExist(
    std::shared_ptr<filesystems::FileSystem>& localFs,
    const std::string& folderPath) {
  if (localFs->exists(folderPath)) {
    localFs->rmdir(folderPath);
  }
  EXPECT_FALSE(localFs->exists(folderPath));
}

} // namespace

class ExpressionVerifierUnitTest : public testing::Test, public VectorBuilder {
 public:
  ExpressionVerifierUnitTest() {
    parse::registerTypeResolver();
  }

 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
  }

  core::TypedExprPtr parse_expression(
      const std::string& text,
      const RowTypePtr& rowType) {
    parse::ParseOptions options;
    auto untyped = parse::parseExpr(text, options);
    return core::Expressions::inferTypes(untyped, rowType, pool_.get());
  }

  std::shared_ptr<memory::MemoryPool> pool_{
      memory::memoryManager()->addLeafPool()};
  std::shared_ptr<core::QueryCtx> queryCtx_{core::QueryCtx::create()};
  core::ExecCtx execCtx_{pool_.get(), queryCtx_.get()};
};

TEST_F(ExpressionVerifierUnitTest, persistReproInfo) {
  filesystems::registerLocalFileSystem();
  auto reproFolder = exec::test::TempDirectoryPath::create();
  const auto reproPath = reproFolder->getPath();
  auto localFs = filesystems::getFileSystem(reproPath, nullptr);

  ExpressionVerifierOptions options{false, reproPath.c_str(), false};
  ExpressionVerifier verifier{&execCtx_, options, nullptr};

  auto testReproPersistency = [this](
                                  ExpressionVerifier& verifier,
                                  const std::string& reproPath,
                                  auto& localFs) {
    auto data = make_row_vector({make_flat_vector<int32_t>({1, 2, 3})});
    auto plan = parse_expression("always_throws(c0)", as_row_type(data->type()));

    removeDirecrtoryIfExist(localFs, reproPath);
    POLLUX_ASSERT_THROW(
        verifier.verify(
            {plan}, {{data, SelectivityVector(data->size())}}, nullptr, false),
        "");
    EXPECT_TRUE(localFs->exists(reproPath));
    EXPECT_FALSE(localFs->list(reproPath).empty());
    removeDirecrtoryIfExist(localFs, reproPath);
  };

  // User errors.
  {
    register_function<AlwaysThrowsUserErrorFunction, int32_t, int32_t>(
        {"always_throws"});
    testReproPersistency(verifier, reproPath, localFs);
  }

  // Runtime errors.
  {
    register_function<AlwaysThrowsRuntimeErrorFunction, int32_t, int32_t>(
        {"always_throws"});
    testReproPersistency(verifier, reproPath, localFs);
  }
}

TEST_F(ExpressionVerifierUnitTest, subsetOfRowsToVerify) {
  // Verify that rows passed to verify() are the only ones evaluated.
  ExpressionVerifierOptions options;
  ExpressionVerifier verifier{&execCtx_, options, nullptr};
  {
    register_function<ThrowsRuntimeErrorOnEvenInputFunction, int32_t, int32_t>(
        {"always_throws"});
    auto data = make_row_vector({make_flat_vector<int32_t>({1, 2, 3})});
    auto plan = parse_expression("always_throws(c0)", as_row_type(data->type()));
    POLLUX_ASSERT_THROW(
        verifier.verify(
            {plan}, {{data, SelectivityVector(data->size())}}, nullptr, false),
        "");

    SelectivityVector rows(data->size());
    rows.setValid(1, false);
    rows.updateBounds();
    verifier.verify({plan}, {{data, rows}}, nullptr, false);
  }
}

} // namespace kumo::pollux::test
